import numpy as np
from cv2 import cv2
import re
import matplotlib.pyplot as plt
import random
from sklearn.cluster import KMeans
from sklearn import decomposition
import matplotlib
import matplotlib.image as gImage
from sklearn.manifold import TSNE
from matplotlib.ticker import FuncFormatter
import scipy.stats
import time
import random
from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report
import os
import math
from scipy.cluster.hierarchy import dendrogram


def calculateDis(GNSS,posA,posB):# posA、B是两个位置的索引，返回这两个位置之间的直线距离和累计距离
    res1 = np.sqrt(np.square(GNSS[posA][0]-GNSS[posB][0])+ np.square(GNSS[posA][1]-GNSS[posB][1]))
    res2 = 0
    for i in range(posA+1,posB+1):
        res2 += np.sqrt(np.square(GNSS[i][0]-GNSS[i-1][0])+ np.square(GNSS[i][1]-GNSS[i-1][1]))
    return res1, res2 

def calculateAngel(GNSS,posA,posB):# # posA、B是两个位置的索引，返回这两个位置之间的方向差，范围为[0,180]
    lenth = len(GNSS) - 1
    if posA <= 0 or posA >= lenth or posB <= 0 or posB >= lenth:
        return 0
    angleA = GNSS[posA][2]
    angleB = GNSS[posB][2]
    delta = abs(angleA - angleB)
    return 360 - delta if delta > 180 else delta


def getAnchorPosNegIdx(path): # path是数据集路径，返回的是每个锚点及其对应正/负样本的id，数据组织形式是[[anchor1,[pos1,pos2,...]],[anchor2,[pos1,pos2,...]]...]
    dirs = []
    for _, dirs, _ in os.walk(path):
        break
    dirs = [int(i) for i in dirs]
    dirs = sorted(dirs)
    anchor_pos_list = []
    for d in dirs:
        cur_anchor_pos_list = []
        cur_anchor_pos_list.append(d)
        files = []
        for _, _, files in os.walk(os.path.join(path, str(d))):
            break
        pos_list = []
        for f in files:
            match = re.match('(\d*)\.png',f)
            f = int(match.group(1))
            pos_list.append(f)
        cur_anchor_pos_list.append(pos_list)
        anchor_pos_list.append(cur_anchor_pos_list)
    
    anchor_neg_list = []
    anchor_num = len(anchor_pos_list)
    for i in range(anchor_num):
        cur_anchor_neg_list = [anchor_pos_list[i][0]]
        neg_list = []
        if i == 0:
            neg_list = anchor_pos_list[1][1] + anchor_pos_list[anchor_num-1][1]
        elif i == anchor_num - 1:
            neg_list = anchor_pos_list[anchor_num - 2][1] + anchor_pos_list[0][1]
        else:
            neg_list = anchor_pos_list[i-1][1] + anchor_pos_list[i+1][1]
        cur_anchor_neg_list.append(neg_list)
        anchor_neg_list.append(cur_anchor_neg_list)

    return anchor_pos_list, anchor_neg_list

def getAnchorPosNegIdx2(path, sampleNum):
    files = []
    for _,_,files in os.walk(path):
        pass
    imgNamesInt = [int(i[0:10]) for i in files]
    imgNamesInt = sorted(imgNamesInt)
    anchor_pos_list = []
    # sampleNum = 16
    for i, _ in enumerate(imgNamesInt):
        if i == len(imgNamesInt) - 1:
            break
        cur_anchor_pos_list = []
        cur_anchor_pos_list.append(imgNamesInt[i])
        pos_list = list(np.linspace(imgNamesInt[i], imgNamesInt[i+1], sampleNum + 1).astype('int32'))[0:-1]
        pos_list = list(set(pos_list))
        pos_list.sort()
        cur_anchor_pos_list.append(pos_list)
        anchor_pos_list.append(cur_anchor_pos_list)
    
    # cur_anchor_pos_list = []
    # cur_anchor_pos_list.append(imgNamesInt[-1])
    # tmp_idx = list(range(imgNamesInt[-1], frame_count)) + list(range(0, imgNamesInt[0]))
    # pos_list = tmp_idx[0:len(tmp_idx):int(len(tmp_idx)/sampleNum)]
    # cur_anchor_pos_list.append(pos_list)
    # anchor_pos_list.append(cur_anchor_pos_list)


    anchor_neg_list = []
    anchor_num = len(anchor_pos_list)
    for i in range(anchor_num):
        cur_anchor_neg_list = [anchor_pos_list[i][0]]
        neg_list = []
        if i == 0:
            neg_list = anchor_pos_list[1][1] + anchor_pos_list[anchor_num-1][1]
        elif i == anchor_num - 1:
            neg_list = anchor_pos_list[anchor_num - 2][1] + anchor_pos_list[0][1]
        else:
            neg_list = anchor_pos_list[i-1][1] + anchor_pos_list[i+1][1]
        cur_anchor_neg_list.append(neg_list)
        anchor_neg_list.append(cur_anchor_neg_list)
    return anchor_pos_list, anchor_neg_list

def getAnchorPosNegIdx3(path, sampleNum, sample_interval):
    files = []
    for _,_,files in os.walk(path):
        pass
    imgNamesInt = [int(i[0:10]) for i in files]
    imgNamesInt = sorted(imgNamesInt)
    anchor_pos_list = []
    # sampleNum = 16
    for i, _ in enumerate(imgNamesInt):
        cur_anchor_pos_list = []
        cur_anchor_pos_list.append(imgNamesInt[i])
        pos_list = [imgNamesInt[i]-j*sample_interval for j in range(int(sampleNum/2))] + [imgNamesInt[i]+(j+1)*sample_interval for j in range(int(sampleNum/2))]
        pos_list = list(set(pos_list))
        pos_list.sort()
        cur_anchor_pos_list.append(pos_list)
        anchor_pos_list.append(cur_anchor_pos_list)

    anchor_neg_list = []
    anchor_num = len(anchor_pos_list)
    for i in range(anchor_num):
        cur_anchor_neg_list = [anchor_pos_list[i][0]]
        neg_list = []
        if i == 0:
            neg_list = anchor_pos_list[1][1] + anchor_pos_list[anchor_num-1][1]
        elif i == anchor_num - 1:
            neg_list = anchor_pos_list[anchor_num - 2][1] + anchor_pos_list[0][1]
        else:
            neg_list = anchor_pos_list[i-1][1] + anchor_pos_list[i+1][1]
        cur_anchor_neg_list.append(neg_list)
        anchor_neg_list.append(cur_anchor_neg_list)
    return anchor_pos_list, anchor_neg_list

def getAnchorPosNegIdx4(path): # path是个txt文件路径，该文件中每一行有两个以空格为分割的索引，其中第一个索引表示锚点索引，第二个索引表示该锚点下对应的一个样本
    f = open(path,'r')
    lines = f.readlines()
    f.close()
    anchor_sample_dict = dict()
    for line in lines:
        anchor = int(line.split()[0])
        sample = int(line.split()[1])
        if anchor not in anchor_sample_dict.keys():
            anchor_sample_dict[anchor] = [sample]
        else:
            anchor_sample_dict[anchor].append(sample)

    anchor_pos_list = []
    # sampleNum = 16
    for anchor in anchor_sample_dict.keys():
        cur_anchor_pos_list = []
        cur_anchor_pos_list.append(anchor)
        pos_list = anchor_sample_dict[anchor]
        pos_list.sort()
        cur_anchor_pos_list.append(pos_list)
        anchor_pos_list.append(cur_anchor_pos_list)

    anchor_pos_list.sort(key = lambda x:x[0])
    anchor_neg_list = []
    anchor_num = len(anchor_pos_list)
    for i in range(anchor_num):
        cur_anchor_neg_list = [anchor_pos_list[i][0]]
        neg_list = []
        if i == 0:
            neg_list = anchor_pos_list[1][1] + anchor_pos_list[anchor_num-1][1]
        elif i == anchor_num - 1:
            neg_list = anchor_pos_list[anchor_num - 2][1] + anchor_pos_list[0][1]
        else:
            neg_list = anchor_pos_list[i-1][1] + anchor_pos_list[i+1][1]
        cur_anchor_neg_list.append(neg_list)
        anchor_neg_list.append(cur_anchor_neg_list)
    return anchor_pos_list, anchor_neg_list

def JS_D(p,q):
    M=(p+q)/2
    return 0.5*scipy.stats.entropy(p, M)+0.5*scipy.stats.entropy(q, M)

def inverse_euclidean_distance(x,y):
    res = np.sqrt(np.sum(np.power(x-y,2)))
    return 1/( res)

def cosine_distance(x,y):
    x_len = np.sqrt(np.sum(np.power(x,2)))
    y_len = np.sqrt(np.sum(np.power(y,2)))
    # if abs(x_len - 1) > 1e-6 or abs(y_len - 1) > 1e-6:
    #     raise ValueError('Error! x or/and y are not normalized vector, program exit.')
    return np.sum(x * y)/(x_len * y_len)

def my_softmax(x):
    x = x - np.max(x)
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x)

def my_entropy(probs):
    if abs(np.sum(probs) - 1) > 1e-6:
        raise ValueError('Error! The sum of prob is not 1, but %.3f'%np.sum(probs))
    res = 0
    for i in probs:
        if i != 0:
            res -= i * math.log(i)
    return res

def plot_confusion_matrix(cm, classes, title='Confusion Matrix'):
    plt.figure()

    # plt.figure(figsize=(12, 8), dpi=100)
    # np.set_printoptions(precision=2)

    size = len(classes)
    acc_matrix = np.zeros_like(cm, dtype=float)
    for i in range(size):
        sums = np.sum(cm[i],dtype=float)
        for j in range(size):
            acc_matrix[i][j] = (cm[i][j] / sums)*100.0

    # 在混淆矩阵中每格的概率值
    ind_array = np.arange(len(classes))
    x, y = np.meshgrid(ind_array, ind_array)
    for x_val, y_val in zip(x.flatten(), y.flatten()):
        c = cm[y_val][x_val]
        acc = acc_matrix[y_val][x_val]
        plt.text(x_val, y_val, "%d\n(%.1f%%)" % (c,acc), color='red', va='center', ha='center')
    

        # offset the tick
    tick_marks = np.array(range(len(classes))) + 0.5
    plt.gca().set_xticks(tick_marks, minor=True)
    plt.gca().set_yticks(tick_marks, minor=True)
    plt.gca().xaxis.set_ticks_position('none')
    plt.gca().yaxis.set_ticks_position('none')
    plt.grid(True, which='minor', linestyle='-')
    # plt.gcf().subplots_adjust(bottom=0.15)
    
    xlocations = range(len(classes))
    plt.xticks(xlocations, classes)
    plt.yticks(xlocations, classes)
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.binary)
    plt.title(title)
    plt.colorbar()

    plt.ylabel('Actual label')
    plt.xlabel('Predict label')
    
    

    
    # show confusion matrix
    # plt.show()



def fig2data(fig):
    """
    fig = plt.figure()
    image = fig2data(fig)
    @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
    @param fig a matplotlib figure
    @return a numpy 3D array of RGBA values
    """
    import PIL.Image as Image
    # draw the renderer
    fig.canvas.draw()
 
    # Get the RGBA buffer from the figure
    w, h = fig.canvas.get_width_height()
    buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8)
    buf.shape = (w, h, 4)
 
    # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
    buf = np.roll(buf, 3, axis=2)
    image = Image.frombytes("RGBA", (w, h), buf.tostring())
    image = np.asarray(image)

    return image


def get_peaks_troughs(x, min_value):
    # 有三个基本假设：1. 波峰波谷必须交替出现 2.波谷值是其相邻两波峰之间的最小值，波峰值是其相邻两波谷之间的最大值 3. 任意两个相邻的波峰波谷值的差必须大于阈值min_value
    res = list()
    data_size = len(x)
    # 1 for peak
    # 2 for through
    last_state = 0
    for i in range(1, data_size - 1):
        state = 0
        if x[i] > x[i-1] and x[i] > x[i+1]:
            state = 1
        elif x[i] < x[i-1] and x[i] < x[i+1]:
            state = 2
        else:
            continue
        if not last_state:
            res.append(i)
            last_state = state
            continue

        last_idx = res[-1]
        if last_state == state: # 保证波谷值是其相邻两波峰之间的最小值，波峰值是其相邻两波谷之间的最大值
            if state == 1:
                if x[i] > x[last_idx]:
                    res[-1] = i
            elif state == 2:
                if x[i] < x[last_idx]:
                    res[-1] = i
        
        else: # 保证波峰波谷交替出现，
            if abs(x[i]-x[last_idx]) > min_value: # 保证相邻波峰波谷值的差大于某阈值
                last_state = state
                res.append(i)
    return res 


def plot_dendrogram(model, **kwargs):
    # Create linkage matrix and then plot the dendrogram

    # create the counts of samples under each node
    counts = np.zeros(model.children_.shape[0])
    n_samples = len(model.labels_)
    for i, merge in enumerate(model.children_):
        current_count = 0
        for child_idx in merge:
            if child_idx < n_samples:
                current_count += 1  # leaf node
            else:
                current_count += counts[child_idx - n_samples]
        counts[i] = current_count

    linkage_matrix = np.column_stack([model.children_, model.distances_,
                                      counts]).astype(float)

    # Plot the corresponding dendrogram
    dendrogram(linkage_matrix, **kwargs)


if __name__ == '__main__':
    posFile = 'D:\\Research\\2020ContrastiveLearningForSceneLabel\\Data\\campus_img_dataset\\pos.npy' # 把字幕文件里的经纬度提取出来保存成numpy数据，这样方便读取节约时间
    GNSS = np.load(posFile)
    dis1,dis2 = calculateDis(GNSS, 1200,1800)
    print(dis1,dis2)

    datasetPath = 'D:\\Research\\2020ContrastiveLearningForSceneLabel\\Data\\campus_img_dataset\\labeledData\\test2\\train'
    anchor_pos_list = getAnchorPosNegIdx(datasetPath)

    print()

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count